// Dustin Escoffery
// Copyright 2007-2010
// Portable Executable Helper API
#include "pehlpapi.h"
#include <windows.h>

// Gets the import descriptor for hModule
PIMAGE_IMPORT_DESCRIPTOR GetImportDescriptor(HMODULE hModule)
{
   PCHAR pImageBase =
      (PCHAR)hModule;
   PIMAGE_DOS_HEADER pDOS =
      PIMAGE_DOS_HEADER(pImageBase);
   PIMAGE_NT_HEADERS pNT =
      (PIMAGE_NT_HEADERS)(pImageBase + pDOS->e_lfanew);
   IMAGE_DATA_DIRECTORY imgDir =
      pNT->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
   DWORD dwImportDirectory =
      imgDir.VirtualAddress;
   return (PIMAGE_IMPORT_DESCRIPTOR)(pImageBase + dwImportDirectory);
}

// Hooks an import of hModule
bool WriteImportAddress(HMODULE hModule, PCSTR pszModuleName,
   LPCSTR lpProcName, PROC pfnHooked)
{
   // Get address of the target function
   PROC pfnCurrent = GetProcAddress(
      GetModuleHandle(pszModuleName), lpProcName);

   // Get the description of the import section
   PIMAGE_IMPORT_DESCRIPTOR pImportDesc = GetImportDescriptor(hModule);

   // Find target DLL
   while (pImportDesc->Name != 0)
   {
      PSTR pszModName = (PSTR)((PBYTE)hModule + pImportDesc->Name);
      if (_stricmp(pszModName, pszModuleName) == 0)
      {
         // Get import address table
         PIMAGE_THUNK_DATA pThunk = (PIMAGE_THUNK_DATA)
            ((PBYTE)hModule + pImportDesc->FirstThunk);

         // Find target import
         while (pThunk->u1.Function != 0)
         {
            PROC* ppfn = (PROC*)&pThunk->u1.Function;
            if (*ppfn == pfnCurrent)
            {
               // Set writable permission
               MEMORY_BASIC_INFORMATION mbi;
               VirtualQuery(ppfn, &mbi,
                  sizeof(MEMORY_BASIC_INFORMATION));
               VirtualProtect(mbi.BaseAddress, mbi.RegionSize,
                  PAGE_READWRITE, &mbi.Protect);

               // Hook the function
               *ppfn = *pfnHooked;

               // Restore permissions
               DWORD dwOldProtect;
               VirtualProtect(mbi.BaseAddress, mbi.RegionSize,
                  mbi.Protect, &dwOldProtect);
               return true;
            }
            pThunk++;
         }
      }
      pImportDesc++;
   }
   return false;
}
